################################################################################
# Libraries and Import
################################################################################
rm(list=ls())

# Libraries
library(sandwich)
library(tidyverse)
library(ri2)
library(stargazer)

# Import
load(file="data/dat_analysis_JP.RData")
jp <- vignette
load(file="data/dat_analysis_CN.RData")
cn <- dat_vignette

# Set seed
set.seed(30)

# Control number of simulations in randomization procedures
sims = 10000 # Set to 10,000 for replication, set to 10 now for speed

################################################################################
# Split each country sample into two datasets by treatment group
################################################################################
# Split dataset by treatment group - Japan
nat_jp = jp %>% filter(!is.na(nat)) %>% rename(z = nat)
econ_jp = jp %>% filter(!is.na(econ)) %>% rename(z = econ)

# Convert China response variable from factor to numeric
cn$D1 = as.numeric(as.character(cn$D1))

# Split dataset by treatment group - China
nat_cn = cn %>% filter(vignette_group == 1 | vignette_group == 2)
econ_cn = cn %>% filter(vignette_group == 3 | vignette_group == 4)

nat_cn$z = with(nat_cn, ifelse(vignette_group == 1, 0, 1))
econ_cn$z = with(econ_cn, ifelse(vignette_group == 3, 0, 1))

################################################################################
# Test for heterogenous effects by country
################################################################################
# Nationalist: Calculate difference in CATEs by country
nat_japan = lm(D1 ~ z, data = nat_jp)
nat_china = lm(D1 ~ z, data = nat_cn)
dic <- abs(nat_japan$coefficients[[2]] - nat_china$coefficients[[2]])

# Use randomization inference to test null that CATEs in both groups equal ATE
nat_jp_reduced = nat_jp %>%
  mutate(country = "Japan") %>%
  select("country", "D1", "z")

nat_cn_reduced = nat_cn %>%
  mutate(country = "China") %>%
  select("country", "D1", "z")

nat_jp_cn = rbind(nat_jp_reduced, nat_cn_reduced)

null <- nat_jp_cn
dic_null <- vector(mode = "integer", length = sims)

# 1. # Form full schedule of potential outcomes assuming constant effects
null$ate = lm(D1 ~ z, data = null)$coefficients[2]
null$D1_null = with(null, ifelse(z == 1, D1, D1 + ate))

for (i in seq_along(dic_null)) { 
  
  # 2. Assign subjects randomly to treatment or control
  null$z_null = complete_ra(nrow(null))
  
  # 3. Calculate difference between estimated CATEs
  nat_null_1 = lm(D1_null ~ z_null, data = null, subset = country == "Japan")
  nat_null_2 = lm(D1_null ~ z_null, data = null, subset = country == "China")
  dic_null[[i]] <- abs(nat_null_1$coefficients[[2]] - nat_null_2$coefficients[[2]])
}

# 4. Calculate two sided p-value
p_nat_country <- sum(abs(dic_null) >= dic)/length(dic_null)

# Economic: Calculate difference in CATEs by country
econ_japan = lm(D1 ~ z, data = econ_jp)
econ_china = lm(D1 ~ z, data = econ_cn)
dic <- abs(econ_japan$coefficients[[2]] - econ_china$coefficients[[2]])

# Use randomization inference to test null that CATEs in both groups equal ATE
econ_jp_reduced = econ_jp %>%
  mutate(country = "Japan") %>%
  select("country", "D1", "z")

econ_cn_reduced = econ_cn %>%
  mutate(country = "China") %>%
  select("country", "D1", "z")

econ_jp_cn = rbind(econ_jp_reduced, econ_cn_reduced)

null <- econ_jp_cn
dic_null <- vector(mode = "integer", length = sims)

# 1. # Form full schedule of potential outcomes assuming constant effects
null$ate = lm(D1 ~ z, data = null)$coefficients[2]
null$D1_null = with(null, ifelse(z == 1, D1, D1 + ate))

for (i in seq_along(dic_null)) { 
  
  # 2. Assign subjects randomly to treatment or control
  null$z_null = complete_ra(nrow(null))
  
  # 3. Calculate difference between estimated CATEs
  econ_null_1 = lm(D1_null ~ z_null, data = null, subset = country == "Japan")
  econ_null_2 = lm(D1_null ~ z_null, data = null, subset = country == "China")
  dic_null[[i]] <- abs(econ_null_1$coefficients[[2]] - econ_null_2$coefficients[[2]])
}

# 4. Calculate two sided p-value
p_econ_country <- sum(abs(dic_null) >= dic)/length(dic_null)

###############################################################################
# Charts of heterogenous effects: All Parties
###############################################################################
groupmeans_nat_jp <- nat_jp %>%
  group_by(z) %>%
  summarise(groupmean = mean(D1, na.rm = TRUE), 
            size = n(),
            var = var(D1)) %>%
  mutate(country = "Japan")

groupmeans_nat_cn <- nat_cn %>%
  group_by(z) %>%
  summarise(groupmean = mean(D1, na.rm = TRUE), 
            size = n(),
            var = var(D1)) %>%
  mutate(country = "China")

# Calculate standard errors
groupmeans_nat_jp$se <- sqrt(groupmeans_nat_jp$var/groupmeans_nat_jp$size)
groupmeans_nat_cn$se <- sqrt(groupmeans_nat_cn$var/groupmeans_nat_cn$size)

# Combine groupmeans estimates
groupmeans = rbind(groupmeans_nat_jp, groupmeans_nat_cn)

# Create chart
ate <- ggplot(groupmeans, aes(factor(z), groupmean, fill = factor(z))) +
  geom_col(position = "dodge", 
           color = "black", alpha = 0.75,
           width = 0.65) + 
  scale_fill_manual(values = c("steelblue2", "seagreen3"),
                    guide = FALSE) +
  geom_errorbar(aes(ymin = groupmean - 1.96*se, ymax = groupmean + 1.96*se),
                position = position_dodge(),
                color="black", alpha = 0.5,
                size=0.5,
                width = 0.65) +
  geom_text(aes(label = round(groupmean, 2)), nudge_y = 0.25, size = 3) +
  facet_wrap( ~ country, ncol = 2) + 
  geom_hline(yintercept = 0, linetype = "dashed") +
  xlab("Treatment") +
  ylab("Mean approval of government's action") + 
  scale_x_discrete("", labels = c("0" = "No interference", "1" = "Interference")) +
  scale_y_continuous(limits = c(0, 5)) +
  theme_classic() +
  theme(plot.title = element_text(hjust = 0.5)) +
  theme(axis.text=element_text(size = 10)) +
  theme(legend.position = "none")

ggsave("figs/5. ate_china_japan_nationalism.pdf", ate, height = 5, width = 6)

# Economic Repression
groupmeans_econ_jp <- econ_jp %>%
  group_by(z) %>%
  summarise(groupmean = mean(D1, na.rm = TRUE), 
            size = n(),
            var = var(D1)) %>%
  mutate(country = "Japan")

groupmeans_econ_cn <- econ_cn %>%
  group_by(z) %>%
  summarise(groupmean = mean(D1, na.rm = TRUE), 
            size = n(),
            var = var(D1)) %>%
  mutate(country = "China")

# Calculate standard errors
groupmeans_econ_jp$se <- sqrt(groupmeans_econ_jp$var/groupmeans_econ_jp$size)
groupmeans_econ_cn$se <- sqrt(groupmeans_econ_cn$var/groupmeans_econ_cn$size)

# Combine groupmeans estimates
groupmeans = rbind(groupmeans_econ_jp, groupmeans_econ_cn)

# Create chart
ate <- ggplot(groupmeans, aes(factor(z), groupmean, fill = factor(z))) +
  geom_col(position = "dodge", 
           color = "black", alpha = 0.75,
           width = 0.65) + 
  scale_fill_manual(values = c("steelblue2", "seagreen3"),
                    guide = FALSE) +
  geom_errorbar(aes(ymin = groupmean - 1.96*se, ymax = groupmean + 1.96*se),
                position = position_dodge(),
                color="black", alpha = 0.5,
                size=0.5,
                width = 0.65) +
  geom_text(aes(label = round(groupmean, 2)), nudge_y = 0.25, size = 3) +
  facet_wrap( ~ country, ncol = 2) + 
  geom_hline(yintercept = 0, linetype = "dashed") +
  xlab("Treatment") +
  ylab("Mean approval of government's action") + 
  scale_x_discrete("", labels = c("0" = "No interference", "1" = "Interference")) +
  scale_y_continuous(limits = c(0, 5)) +
  theme_classic() +
  theme(plot.title = element_text(hjust = 0.5)) +
  theme(axis.text=element_text(size = 10)) +
  theme(legend.position = "bottom")

ggsave("figs/6. ate_china_japan_economic.pdf", ate, height = 5, width = 6)
